import torch
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import LinearSegmentedColormap

# 示例矩阵数据
data = np.zeros([5,5])  # 生成一个 10x10 的随机矩阵，值在 [0, 1] 之间
def generate_head_map(data_tensor, save_path=None):
    print("hello world")
    data=data_tensor.detach().clone()
    data=data.to("cpu")
    data=data.numpy()
    # 检查数据范围
    assert np.min(data) >= 0 and np.max(data) <= 1, "数据范围必须在 [0, 1] 之间"

    # 自定义颜色映射
    # 0 对应无色（透明），1 对应深红色
    colors = [(1, 1, 1, 0), (1, 0, 0, 1)]  # (R, G, B, alpha)
    cmap_name = 'custom_red'
    custom_cmap = LinearSegmentedColormap.from_list(cmap_name, colors)

    # 绘制热力图
    plt.figure(figsize=(8, 6))  # 设置图像大小
    heatmap = plt.imshow(data, cmap=custom_cmap, aspect='auto', vmin=0, vmax=1)  # 使用自定义颜色映射
    plt.colorbar(heatmap)  # 添加颜色条

    # 添加标题和标签
    plt.title('Heatmap Example with Custom Red Color Scheme')
    plt.xlabel('X-axis')
    plt.ylabel('Y-axis')

    # 保存图像
    if save_path is not None:
        plt.savefig(save_path)  # 保存为 PNG 文件
        print(f'image saved as {save_path}')
    else:
        print("image not saved")

    # 显示图像（可选）
    plt.show()